#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
手动创建大乐透和双色球基础LSTM模型文件的脚本
"""

import os
import sys
import torch
import torch.nn as nn
from datetime import datetime

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.append(project_root)

class SimpleLstmCRFModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, output_seq_length, num_layers=1):
        super(SimpleLstmCRFModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.output_seq_length = output_seq_length
        
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_seq_length * output_dim)
        # 简化版本，不使用CRF层，因为可能缺少依赖
        
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        fc_out = self.fc(lstm_out[:, -1, :])
        fc_out = fc_out.view(-1, self.output_seq_length, self.output_dim)
        return fc_out

def create_dlt_model():
    """创建大乐透基础LSTM模型"""
    print("开始创建大乐透基础LSTM模型...")
    
    # 配置路径
    model_save_path = os.path.join(project_root, 'scripts', 'dlt', 'dlt_lstm_model.pth')
    print(f"模型保存路径: {model_save_path}")
    
    try:
        # 创建目录（如果不存在）
        os.makedirs(os.path.dirname(model_save_path), exist_ok=True)
        
        # 初始化模型参数
        input_dim = 7
        hidden_dim = 128
        red_output_dim = 35
        blue_output_dim = 12
        red_output_seq_length = 5
        blue_output_seq_length = 2
        
        # 创建模型实例
        red_model = SimpleLstmCRFModel(input_dim, hidden_dim, red_output_dim, red_output_seq_length)
        blue_model = SimpleLstmCRFModel(input_dim, hidden_dim, blue_output_dim, blue_output_seq_length)
        
        # 初始化权重
        def init_weights(m):
            if isinstance(m, nn.Linear):
                torch.nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(0.01)
            elif isinstance(m, nn.LSTM):
                for name, param in m.named_parameters():
                    if 'weight' in name:
                        torch.nn.init.xavier_uniform_(param)
                    elif 'bias' in name:
                        param.data.fill_(0)
        
        red_model.apply(init_weights)
        blue_model.apply(init_weights)
        
        # 保存模型状态字典
        torch.save({
            'red_model': red_model.state_dict(),
            'blue_model': blue_model.state_dict(),
            'input_dim': input_dim,
            'hidden_dim': hidden_dim,
            'red_output_dim': red_output_dim,
            'blue_output_dim': blue_output_dim,
            'red_output_seq_length': red_output_seq_length,
            'blue_output_seq_length': blue_output_seq_length,
            'created_time': datetime.now().isoformat()
        }, model_save_path)
        
        print("大乐透基础LSTM模型创建成功！")
        print(f"模型文件已保存到: {model_save_path}")
        return True
        
    except Exception as e:
        print(f"创建大乐透模型过程出错: {e}")
        import traceback
        print(traceback.format_exc())
        return False

def create_ssq_model():
    """创建双色球基础LSTM模型"""
    print("开始创建双色球基础LSTM模型...")
    
    # 配置路径
    model_save_path = os.path.join(project_root, 'scripts', 'ssq', 'ssq_lstm_model.pth')
    print(f"模型保存路径: {model_save_path}")
    
    try:
        # 创建目录（如果不存在）
        os.makedirs(os.path.dirname(model_save_path), exist_ok=True)
        
        # 初始化模型参数
        input_dim = 7
        hidden_dim = 128
        red_output_dim = 33
        blue_output_dim = 16
        red_output_seq_length = 6
        blue_output_seq_length = 1
        
        # 创建模型实例
        red_model = SimpleLstmCRFModel(input_dim, hidden_dim, red_output_dim, red_output_seq_length)
        blue_model = SimpleLstmCRFModel(input_dim, hidden_dim, blue_output_dim, blue_output_seq_length)
        
        # 初始化权重
        def init_weights(m):
            if isinstance(m, nn.Linear):
                torch.nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(0.01)
            elif isinstance(m, nn.LSTM):
                for name, param in m.named_parameters():
                    if 'weight' in name:
                        torch.nn.init.xavier_uniform_(param)
                    elif 'bias' in name:
                        param.data.fill_(0)
        
        red_model.apply(init_weights)
        blue_model.apply(init_weights)
        
        # 保存模型状态字典
        torch.save({
            'red_model': red_model.state_dict(),
            'blue_model': blue_model.state_dict(),
            'input_dim': input_dim,
            'hidden_dim': hidden_dim,
            'red_output_dim': red_output_dim,
            'blue_output_dim': blue_output_dim,
            'red_output_seq_length': red_output_seq_length,
            'blue_output_seq_length': blue_output_seq_length,
            'created_time': datetime.now().isoformat()
        }, model_save_path)
        
        print("双色球基础LSTM模型创建成功！")
        print(f"模型文件已保存到: {model_save_path}")
        return True
        
    except Exception as e:
        print(f"创建双色球模型过程出错: {e}")
        import traceback
        print(traceback.format_exc())
        return False

def main():
    """主函数"""
    print("彩票基础LSTM模型创建程序")
    print(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print()
    
    # 创建大乐透模型
    print("=" * 50)
    dlt_success = create_dlt_model()
    
    print()
    
    # 创建双色球模型
    print("=" * 50)
    ssq_success = create_ssq_model()
    
    print()
    print("=" * 50)
    if dlt_success and ssq_success:
        print("所有模型创建完成！")
    else:
        print("部分模型创建失败！")
    
    print(f"结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

if __name__ == "__main__":
    main()